﻿#include "SGHSystem.h"
#include "Application.h"
#include "Mesh.h"
#include "Timer.h"
#include "Transient.h"

#include "Algorithm.h"
#include <vtkm/cont/ArrayHandleGroupVec.h>
#include <vtkm/worklet/CellMeasure.h>
#include <vtkm/worklet/WorkletMapTopology.h>

//struct ComputeCellAreaAndMass : vtkm::worklet::WorkletVisitCellsWithPoints
//{
//  using ControlSignature =
//    void(CellSetIn, FieldInPoint, FieldInCell density, FieldOut area, FieldOutCell mass);
//  using ExecutionSignature = void(_2, _3, _4, _5);
//  template<typename PType>
//  VTKM_EXEC void operator()(const PType& pts, const Real& density, Real& area, Real& mass) const
//  {
//    typename PType::ComponentType v1 = pts[1] - pts[0];
//    typename PType::ComponentType v2 = pts[2] - pts[0];
//    area = (0.5) * (vtkm::Magnitude(vtkm::Cross(v1, v2)));
//    mass = density * area;
//  }
//};
//
//struct ComputeCellEnergy : vtkm::worklet::WorkletMapField
//{
//  using ControlSignature = void(FieldIn area, FieldIn density, FieldIn pressure, FieldOut energy);
//  using ExecutionSignature = _4(_1, _2, _3);
//  VTKM_EXEC Real operator()(const Real& area, const Real& density, const Real& pressure) const
//  {
//    Real gamma = 5. / 3;
//    Real e = pressure / (gamma - 1) / density;
//    return area * e;
//  };
//};

struct InitCellEnergy : WorkletMapField
{
  using ControlSignature = void(FieldIn density, FieldIn pressure, FieldOut energy);
  using ExecutionSignature = _3(_1, _2);
  VTKM_EXEC Real operator()(const Real& density, const Real& pressure) const
  {
    Real gamma = 5. / 3;
    Real e = pressure / (gamma - 1) / density;
    return e;
  };
};

struct InitCellMass : WorkletMapField
{
  using ControlSignature = void(FieldIn density, FieldIn area, FieldOut mass);
  using ExecutionSignature = _3(_1, _2);
  VTKM_EXEC Real operator()(const Real& density, const Real& area) const { return density * area; };
};

struct InitPointMass : vtkm::worklet::WorkletVisitPointsWithCells
{
  using ControlSignature = void(CellSetIn, FieldInCell cell_mass, FieldOutPoint point_mass);
  using ExecutionSignature = _3(CellCount, _2);

  template<typename CellValueVecType>
  VTKM_EXEC Real operator()(const IdComponent& cell_count, const CellValueVecType& cell_mass) const
  {
    Real point_mass = 0;
    for (IdComponent i = 0; i < cell_count; i++)
    {
      point_mass += cell_mass[i];
    }
    point_mass /= 3.;
    return point_mass;
  }
};
struct ComputeCellToPointForce : WorkletVisitCellsWithPoints
{
  // clang-format off
  using ControlSignature =  void(CellSetIn, FieldInPoint, FieldInCell pressure, FieldOut force, FieldOut c2n);
  // clang-format on
  using ExecutionSignature = void(_2, _3, _4, PointIndices, _5);

  template<typename PType, typename ForceType, typename C2NType>
  VTKM_EXEC void operator()(const PType& pts,
                            const Real& pressure,
                            ForceType& force,
                            const C2NType& c2n,
                            Vec<Id, 3>& conn) const
  {
    Vec3f normal;
    normal = Normal(pts[1], pts[2]);
    //force[0] = { 0, 0, 0 };
    force[0] = pressure * normal;

    normal = Normal(pts[2], pts[0]);
    //force[1] = { 0, 0, 0 };
    force[1] = pressure * normal;

    normal = Normal(pts[0], pts[1]);
    //force[2] = { 0, 0, 0 };
    force[2] = pressure * normal;

    conn[0] = c2n[0];
    conn[1] = c2n[1];
    conn[2] = c2n[2];
  }

  VTKM_EXEC Vec3f Normal(const Vec3f& ps, const Vec3f& pe) const
  {
    return Vec3f{ ps[1] - pe[1], pe[0] - ps[0], 0 } / 2;
  }
};


struct PeriodicMap : vtkm::worklet::WorkletMapField
{
  using ControlSignature = void(FieldIn boundary_id, FieldOut periodic_map);
  using ExecutionSignature = void(_1, _2, WorkIndex);

  PeriodicMap(int nx, int ny)
    : _nx(nx)
    , _ny(ny)
  {
  }
  VTKM_EXEC
  void operator()(const Id& bid, Id& periodic_map, Id tid) const
  {
    Vec3f other_p;
    if (bid == 0) // bottom
    {
      periodic_map = tid + (_nx + 1) * _ny;
      //std::cout << tid << " " << periodic_map << std::endl;
    }
    else if (bid == 2) // top
    {
      periodic_map = tid - (_nx + 1) * _ny;
    }
    else if (bid == 3) // left
    {
      periodic_map = tid + _nx;
    }
    else if (bid == 1) // right
    {
      periodic_map = tid - _nx;
    }
    else
    {
      periodic_map = -1;
    }
  }
  int _nx, _ny;
};

struct ComputePointForce : vtkm::worklet::WorkletReduceByKey
{
  struct ValuesInWholeArray : vtkm::cont::arg::ControlSignatureTagBase
  {
    using TypeCheckTag = vtkm::cont::arg::TypeCheckTagArrayIn;
    using TransportTag = vtkm::cont::arg::TransportTagKeyedValuesIn;
    using FetchTag = vtkm::exec::arg::FetchTagExecObject;
  };

  using ControlSignature = void(KeysIn,
                                ValuesIn c2n_force,
                                ReducedValuesOut point_force,
                                ValuesInWholeArray,
                                ReducedValuesIn bid,
                                ReducedValuesIn periodic_map);
  using ExecutionSignature = void(_1, _2, _3, _4, _5, _6, WorkIndex);

  template<typename ValueType, typename WholeArray>
  VTKM_EXEC void operator()(const Id& key,
                            const ValueType& c2n_force,
                            Vec3f& point_force,
                            WholeArray& arr,
                            const Id& bid,
                            const Id& map_id,
                            Id& tid) const

  {
    point_force = Vec3f{ 0, 0, 0 };
    auto& force = arr.Get(key);
    for (IdComponent i = 0; i < force.GetNumberOfComponents(); i++)
    {
      //std::cout << key << " " << force.GetNumberOfComponents() << "  "
      //          << c2n_force.GetNumberOfComponents() << "\n";
      point_force += force[i];
    }

    //std::cout << bid << " " << std::endl;
    if (map_id >= 0)
    {
      auto& force = arr.Get(map_id);
      //std::cout << key << " " << map_id << "\n";
      for (IdComponent i = 0; i < force.GetNumberOfComponents(); i++)
      {
        point_force += force[i];
      }
    }
    if (bid == 4)
    {
      point_force = Vec3f{ 0, 0, 0 };
    }
  }
};

struct UpdateVelocityAndPosition : vtkm::worklet::WorkletMapField
{
  using ControlSignature = void(FieldIn point_mass,
                                FieldIn force,
                                FieldInOut velocity,
                                FieldInOut position);
  using ExecutionSignature = void(_1, _2, _3, _4);

  UpdateVelocityAndPosition(Real dt)
    : _dt(dt)
  {
  }
  template<typename CoordType>
  VTKM_EXEC void operator()(const Real& point_mass,
                            const Vec3f& force,
                            Vec3f& vel,
                            CoordType& point) const
  {
    vel += force * _dt / point_mass;
    point += vel * _dt;
  }

  Real _dt;
};

struct UpdateCellEnergy : WorkletVisitCellsWithPoints
{
  // clang-format off
  using ControlSignature =  void(CellSetIn, FieldInCell mass, FieldInPoint velocity, FieldInCell force, FieldInOutCell energy);
  // clang-format on
  using ExecutionSignature = void(PointCount, _2, _3, _4, _5);

  UpdateCellEnergy(Real dt)
    : _dt(dt)
  {
  }
  template<typename ForceType, typename VelocityType>
  VTKM_EXEC void operator()(const IdComponent& num_pts,
                            const Real& mass,
                            const VelocityType& velocity,
                            const ForceType& force,
                            Real& energy) const
  {
    Real rhs = 0;
    for (IdComponent i = 0; i < num_pts; i++)
    {
      rhs += vtkm::Dot(force[i], velocity[i]);
    }
    energy -= _dt * rhs / mass;
  }

  Real _dt;
};


struct ComputeCellAreaAndDensity : vtkm::worklet::WorkletVisitCellsWithPoints
{
  using ControlSignature =
    void(CellSetIn, FieldInPoint, FieldInCell mass, FieldOut area, FieldOutCell density);
  using ExecutionSignature = _5(_2, _3, _4);
  template<typename PType>
  VTKM_EXEC Real operator()(const PType& pts, const Real& mass, Real& area) const
  {
    typename PType::ComponentType v1 = pts[1] - pts[0];
    typename PType::ComponentType v2 = pts[2] - pts[0];
    area = (0.5) * (vtkm::Magnitude(vtkm::Cross(v1, v2)));
    return mass / area;
  }
};

struct ComputeCellPressure : vtkm::worklet::WorkletMapField
{
  using ControlSignature = void(FieldIn density, FieldIn energy, FieldOut pressure);
  using ExecutionSignature = _3(_1, _2);
  VTKM_EXEC Real operator()(const Real& density, const Real& energy) const
  {
    Real gamma = 5. / 3;
    Real pressure = (gamma - 1) * density * energy;
    return pressure;
  };
};
SGHSystem::SGHSystem(const Configuration& cfg)
  : System(cfg)
{
}

void SGHSystem::Init()
{
  _data_set.SetCellSet(_mesh.GetCellSet());
  _data_set.AddCoordinateSystem(_mesh.GetCoordinate());

  // 通过初始条件的赋值变量
  this->AddCellVariable<Real>("pressure");
  this->AddCellVariable<Real>("density");
  this->AddPointVariable<Vec3f>("velocity");

  // 对 density, velocity, pressure赋值
  this->InitialCondition();

  // AuxVariable，辅助变量
  this->AddCellVariable<Real>("area");
  this->AddCellVariable<Real>("energy");
  this->AddCellVariable<Real>("cell_mass");
  this->AddPointVariable<Real>("point_mass");
  this->AddPointVariable<Vec3f>("force");


  this->InitVariable();

  // 建立周期边界映射
  auto& bnd_id = _mesh.GetBoundaryId();
  Invoker{}(PeriodicMap{ _mesh._nx, _mesh._ny }, bnd_id, _periodic_map);
}

// 初始化
// 变量: energy,
// 辅助变量: area,
// 辅助变量: point_mass, cell_mass
void SGHSystem::InitVariable()
{
  auto& density = GetVariableArray<Real>("density");
  auto& pressure = GetVariableArray<Real>("pressure");
  auto& energy = GetVariableArray<Real>("energy");

  auto& area = GetVariableArray<Real>("area");
  invoker(
    vtkm::worklet::CellMeasure<vtkm::Area>{}, _mesh.GetCellSet(), _mesh.GetCoordinate(), area);

  auto& cell_mass = GetVariableArray<Real>("cell_mass");
  invoker(InitCellEnergy{}, density, pressure, energy);
  invoker(InitCellMass{}, density, area, cell_mass);

  auto& point_mass = GetVariableArray<Real>("point_mass");
  invoker(InitPointMass{}, _mesh.GetCellSet(), cell_mass, point_mass);
}

// 计算 CellToPointForce,
void SGHSystem::PreSolve()
{
  UpdateCellToPointForce();
  UpdatePointForce();
}

void SGHSystem::UpdateCellToPointForce()
{
  auto& pressure = GetVariableArray<Real>("pressure");
  invoker(ComputeCellToPointForce{},
          _mesh.GetCellSet(),
          _mesh.GetCoordinate(),
          pressure,
          vtkm::cont::make_ArrayHandleGroupVec<3>(_c2n_force),
          vtkm::cont::make_ArrayHandleGroupVec<3>(_c2n));
}

void SGHSystem::UpdatePointForce()
{
  auto& pressure = GetVariableArray<Real>("pressure");
  invoker(ComputeCellToPointForce{},
          _mesh.GetCellSet(),
          _mesh.GetCoordinate(),
          pressure,
          vtkm::cont::make_ArrayHandleGroupVec<3>(_c2n_force),
          vtkm::cont::make_ArrayHandleGroupVec<3>(_c2n));

  auto& force = GetVariableArray<Vec3f>("force");
  auto& bid = _mesh.GetBoundaryId();
  vtkm::worklet::Keys<Id> keys(_c2n);
  invoker(ComputePointForce{}, keys, _c2n_force, force, _c2n_force, bid, _periodic_map);
}
void SGHSystem::ComputeTimeStepper()
{
  auto& transient = static_cast<Transient&>(*(_app.GetExecutioner()));
  _dt = transient.Dt();
}

// 时间推进, 更新 velocity, position, energy
void SGHSystem::Solve()
{
  ComputeTimeStepper(); // 计算 dt

  auto point_mass = GetVariableArray<Real>("point_mass");
  auto force = GetVariableArray<Vec3f>("force");
  auto velocity = GetVariableArray<Vec3f>("velocity");
  auto position = _mesh.GetCoordinate();
  invoker(UpdateVelocityAndPosition{ _dt }, point_mass, force, velocity, position);

  auto cell_mass = GetVariableArray<Real>("cell_mass");
  auto energy = GetVariableArray<Real>("energy");

  invoker(UpdateCellEnergy{ _dt },
          _mesh.GetCellSet(),
          cell_mass,
          velocity,
          vtkm::cont::make_ArrayHandleGroupVec<3>(_c2n_force),
          energy);
}

// 计算 area, density, 用来输出流场
// 计算 pressure, 下个PreSolve计算force
void SGHSystem::PostSolve()
{
  auto& cell_mass = GetVariableArray<Real>("cell_mass");
  auto& area = GetVariableArray<Real>("area");
  auto& density = GetVariableArray<Real>("density");
  invoker(ComputeCellAreaAndDensity{},
          _mesh.GetCellSet(),
          _mesh.GetCoordinate(),
          cell_mass,
          area,
          density);


  auto& energy = GetVariableArray<Real>("energy");
  auto& pressure = GetVariableArray<Real>("pressure");
  invoker(ComputeCellPressure{}, density, energy, pressure);
}


void SGHSystem::Evolve()
{


  Timer timer;
  timer.Start();

  PreSolve();
  Solve();
  PostSolve();

  Info("SGHSystem::Evolve运行完成，", "时间: ", timer.GetElapsedTime(), " s");
}
